package com.codehelper.freemarker;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.config.OutputFile;
import com.baomidou.mybatisplus.generator.config.TemplateConfig;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.engine.AbstractTemplateEngine;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;
import com.codehelper.TemplateFileInfo;
import com.sun.istack.internal.NotNull;

import java.io.File;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 支持自定义模板及类的生成
 *
 * @author xieqiang
 * @date 2024/3/26
 * @since java 1.8
 */
public class CustomFreemarkerTemplateEngine extends FreemarkerTemplateEngine {

    private List<TemplateFileInfo> templateFileInfos;


    public CustomFreemarkerTemplateEngine(List<TemplateFileInfo> templateFileInfos) {
        this.templateFileInfos = templateFileInfos;
    }


    protected void outputCustomFile(@NotNull Map<String, String> customFile, @NotNull TableInfo tableInfo, @NotNull Map<String, Object> objectMap) {
        String entityName = tableInfo.getEntityName();
        objectMap.put("classVar", org.springframework.util.StringUtils.uncapitalize(entityName));
        this.templateFileInfos.forEach(templateFileInfo -> {
            objectMap.put(templateFileInfo.templateName + "Package", templateFileInfo.packageName);
            String className = templateFileInfo.getFileName(entityName);
            objectMap.put(templateFileInfo.templateName + "ClassName",className);
        });

        Set<String> baseClassFields= (Set<String>) objectMap.get("BASE_ENTITY_FIELDS");
        if(baseClassFields==null){
            baseClassFields=new HashSet<>();
        }
        Set<String> finalBaseClassFields = baseClassFields;
        objectMap.put("excludeBaseEntityFields",tableInfo.getFields().stream().filter(item-> finalBaseClassFields.contains(item.getName())).collect(Collectors.toList()));

        customFile.forEach((k, v) -> {
            System.out.println("key:"+k);
            boolean need = (boolean) objectMap.getOrDefault("need_" + k, true);
            if (need) {
                TemplateFileInfo templateFileInfo = this.templateFileInfos.stream().filter(item -> item.templateName.equals(k)).findFirst().get();
                String className = templateFileInfo.getFileName(entityName);
                String fileName = templateFileInfo.genFilePath + File.separator + className + templateFileInfo.fileExt;
                System.out.println(fileName);
                this.outputFile(new File(fileName), objectMap, v);

            }
        });
    }


    @Override
    public AbstractTemplateEngine batchOutput() {
        try {
            ConfigBuilder config = this.getConfigBuilder();
            List<TableInfo> tableInfoList = config.getTableInfoList();
            tableInfoList.forEach((tableInfo) -> {
                Map<String, Object> objectMap = this.getObjectMap(config, tableInfo);
                Map<String, Object> injectMap = config.getInjectionConfig().getCustomMap();
                if (!injectMap.isEmpty()) {
                    objectMap.putAll(injectMap);
                }
                Optional.ofNullable(config.getInjectionConfig()).ifPresent((t) -> {
                    t.beforeOutputFile(tableInfo, objectMap);
                    this.outputCustomFile(t.getCustomFile(), tableInfo, objectMap);
                });
            });
            return this;
        } catch (Exception var3) {
            throw new RuntimeException("无法创建文件，请检查配置信息！", var3);
        }
    }

}
